import asyncio
import websockets
import logging
import ssl
from typing import Dict
import socket
from datetime import datetime

# 配置日志显示中文
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)

class WebSocketServer:
    def __init__(self):
        self.clients: Dict[websockets.WebSocketServerProtocol, str] = {}
        self.message_count = 0
        
    async def string_handler(self, *args):
        """处理/string端点的文本消息"""
        try:
            # 兼容不同版本websockets库的参数传递方式
            if len(args) == 2:
                websocket, path = args
            else:
                websocket = args[0]
                path = websocket.request.path if hasattr(websocket, 'request') else '/string'
            
            client_ip = websocket.remote_address[0]
            logging.info(f"[{datetime.now()}] 收到文本连接 路径: {path} 来自: {client_ip}")
        except Exception as e:
            logging.error(f"连接初始化失败: {e}")
            return
        
        self.clients[websocket] = client_ip
        logging.info(f"当前连接数: {len(self.clients)}")
        
        try:
            async for message in websocket:
                logging.info(f"收到文本消息 [{client_ip}]: {message}")
                # 广播消息给所有客户端
                for client in self.clients:
                    if client != websocket:  # 不发送回原客户端
                        try:
                            await client.send(message)
                        except Exception as e:
                            logging.error(f"发送消息失败: {e}")
                            self.clients.pop(client, None)
                
                self.message_count += 1
                if self.message_count > 400000:
                    await websocket.close(code=1001, reason="GoAway")
                    break
                    
        except websockets.exceptions.ConnectionClosed as e:
            logging.info(f"客户端断开: {client_ip} 原因: {e.reason or '正常关闭'}")
        except Exception as e:
            logging.error(f"处理消息错误: {e}")
        finally:
            self.clients.pop(websocket, None)
            logging.info(f"客户端移除: {client_ip} 剩余连接数: {len(self.clients)}")

    async def binary_handler(self, *args):
        """处理/binary端点的二进制消息"""
        # 兼容不同版本websockets库的参数传递方式
        if len(args) == 2:
            websocket, path = args
        else:
            websocket = args[0]
            path = '/binary'  # 默认路径
            
        client_ip = websocket.remote_address[0]
        logging.info(f"[{datetime.now()}] 收到二进制连接 路径: {path or '/'} 来自: {client_ip}")
        
        self.clients[websocket] = client_ip
        logging.info(f"当前连接数: {len(self.clients)}")
        
        try:
            async for message in websocket:
                logging.info(f"收到二进制消息 [{client_ip}]: {message}")
                # 广播消息给所有客户端
                for client in self.clients:
                    if client != websocket:  # 不发送回原客户端
                        try:
                            await client.send(message)
                        except Exception as e:
                            logging.error(f"发送消息失败: {e}")
                            self.clients.pop(client, None)
                
                self.message_count += 1
                if self.message_count > 400000:
                    await websocket.close(code=1001, reason="GoAway")
                    break
                    
        except websockets.exceptions.ConnectionClosed as e:
            logging.info(f"客户端断开: {client_ip} 原因: {e.reason or '正常关闭'}")
        except Exception as e:
            logging.error(f"处理消息错误: {e}")
        finally:
            self.clients.pop(websocket, None)
            logging.info(f"客户端移除: {client_ip} 剩余连接数: {len(self.clients)}")

async def home(request, response):
    """处理根路径请求"""
    response.headers['Content-Type'] = 'application/octet-stream'
    await response.send(b"Hello Binary")

def is_port_available(port: int) -> bool:
    """检查端口是否可用"""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind(('0.0.0.0', port))
            return True
        except OSError:
            return False

async def start_server(port: int = 80, ssl_context: ssl.SSLContext = None):
    """启动WebSocket服务器"""
    if not is_port_available(port):
        new_port = port + 1
        while new_port < 65535 and not is_port_available(new_port):
            new_port += 1
        if new_port >= 65535:
            logging.error(f"找不到可用端口，从{port}开始尝试")
            return
        logging.warning(f"端口{port}被占用，使用端口{new_port}替代")
        port = new_port

    server = WebSocketServer()
    
    try:
        async with websockets.serve(
            server.string_handler,
            "0.0.0.0",
            port,
            ssl=ssl_context,
            ping_interval=None
        ):
            protocol = "wss" if ssl_context else "ws"
            logging.info(f"WebSocket服务器已启动 {protocol}://0.0.0.0:{port}")
            logging.info(f"文本端点: {protocol}://127.0.0.1:{port}/string")
            logging.info(f"二进制端点: {protocol}://127.0.0.1:{port}/binary")
            await asyncio.Future()  # 永久运行
    except Exception as e:
        logging.error(f"启动服务器失败: {e}")

if __name__ == "__main__":
    print("WebSocket服务器(兼容Go版本)")
    print("按Ctrl+C停止服务器\n")
    
    import argparse
    parser = argparse.ArgumentParser(description='WebSocket服务器')
    parser.add_argument('--http-port', type=int, default=8080, help='HTTP端口 (默认: 8080)')
    parser.add_argument('--https-port', type=int, default=8443, help='HTTPS端口 (默认: 8443)')
    parser.add_argument('--cert', type=str, help='SSL证书文件(可选)')
    parser.add_argument('--key', type=str, help='SSL密钥文件(可选)')
    args = parser.parse_args()

    # 创建SSL上下文
    ssl_context = None
    cert_path = args.cert if args.cert else "WebSocketServer/server.crt"
    key_path = args.key if args.key else "WebSocketServer/server.key"
    
    try:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain(cert_path, key_path)
        logging.info(f"SSL证书加载成功: {cert_path}, {key_path}")
    except FileNotFoundError:
        logging.warning(f"证书文件未找到: {cert_path} 或 {key_path}, 仅提供HTTP服务")
    except Exception as e:
        logging.warning(f"SSL证书加载失败: {e}, 仅提供HTTP服务")

    try:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        # 仅启动HTTP服务
        http_server = start_server(args.http_port)
        loop.create_task(http_server)
        
        logging.info(f"WebSocket服务器运行在 ws://0.0.0.0:{args.http_port}")
        logging.info(f"请确保客户端使用正确的URL连接，例如: ws://[服务器IP]:{args.http_port}/string")
        
        loop.run_forever()
    except KeyboardInterrupt:
        logging.info("\n服务器正在停止...")
    except Exception as e:
        logging.error(f"服务器错误: {e}")
    finally:
        loop.close()
